import { FaissStore } from "@langchain/community/vectorstores/faiss";
import { ScoreThresholdRetriever } from "langchain/retrievers/score_threshold";
import { MultiQueryRetriever } from "langchain/retrievers/multi_query";
import { useModel } from "../model";

export const useScoreThresholdRetriever = (vectorStore: FaissStore) => {
    /**
     * ScoreThresholdRetriever 是 LangChain 提供的一个 Retriever，用于根据相似度阈值过滤检索结果。
     * 它的主要作用是：
     * 1. 根据给定的相似度阈值，过滤掉相似度低于该阈值的文档。
     * 2. 返回过滤后的文档列表。
     * 3. 可以用于优化检索结果的质量和效率。
     */
    const retriever = ScoreThresholdRetriever.fromVectorStore(vectorStore, {
        // 相似度阈值
        minSimilarityScore: 0.95,
        // 返回文档数量最大值
        maxK: 5,
        // 步长
        kIncrement: 1,
    });
    return retriever;
};

export  const useMultiQueryRetriever = (vectorStore: FaissStore) => {
    /**
     * MultiQueryRetriever 是 LangChain 提供的一个 Retriever，用于根据多个查询词进行检索。
     * 它的主要作用是：
     它使用 LLM 去将用户的输入改写成多个不同写法，从不同的角度来表达同一个意思，来克服因为关键词或者细微措词导致检索效果差的问题。
     */
    const model = useModel();
    const retriever = MultiQueryRetriever.fromLLM({
        llm: model,
        queryCount: 5,
        verbose: true,
        retriever: vectorStore.asRetriever(3),
    });
    return retriever;
}